# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import os
import sys

import numpy as np
import cv2
    
def parse_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_om_out', type=str,
                        help='the path to load om output file.',default= "/home/bigbigan/output/")
    parser.add_argument('--input_gpu_generated_img', type=str,
                        help='the path to load generated image file of gpu.',default= "/root/Neesky/SFEW_100/Reference/")
    parser.add_argument('--input_pb_generated_img', type=str,
                        help='the path to load generated image file of pb.',default= "/root/Neesky/SFEW_100/Reference/")
    parser.add_argument('--output_om_generated_img', type=str,
                        help='dir where to save generated image file.',default= "/root/Neesky/SFEW_100/Reference/")
    return parser.parse_args(argv)

def main(args):
    # Load om file
    om_out = np.fromfile(args.input_om_out, dtype='float32').reshape([256,32,32,1])
    generated_images = om_out * 250.0
    generated_images = generated_images.astype(np.uint8)   
    # save generated images of off-line inference(om)
    for i in range(256):
        cv2.imwrite(args.output_om_generated_img+'/{}.png'.format(i+1), generated_images[i,...])  
        
    # Load generated images of gpu
    gpu_out = np.fromfile(args.input_gpu_generated_img, dtype='float32').reshape([256,32,32,1])
    # Load generated images of pb
    pb_out = np.fromfile(args.input_pb_generated_img, dtype='float32').reshape([256,32,32,1])
    
    # Record the digital part of the images generated by pb,om and gpu
    pb_number = []
    om_number = []
    gpu_number = []
    for i in range (256): 
        for j in range(32):
            for k in range(32):
                for l in range(1):
                    pixel_pb = pb_out[i][j][k][l]
                    pixel_om = om_out[i][j][k][l]
                    pixel_gpu = gpu_out[i][j][k][l]
                    # eliminate the influence of background color whose pixel value close to zero
                    if (pixel_om > 1e-4 and pixel_pb > 1e-4 and pixel_gpu > 1e-4):  
                        pb_number.append(pixel_pb)
                        om_number.append(pixel_om)
                        gpu_number.append(pixel_gpu)
    pb_number = np.array(pb_number)   
    om_number = np.array(om_number)
    gpu_number = np.array(gpu_number)

    # The mean difference
    average_error_pb_om = np.abs(pb_number-om_number)
    average_error_gpu_om = np.abs(gpu_number-om_number)
    print("the average error of pixel value between pb and om:",np.mean(average_error_pb_om))
    print("the average error of pixel value between gpu and om:",np.mean(average_error_gpu_om))
    
    # The total difference
    total_error_pb_om = np.abs(pb_number-om_number).sum()
    total_error_gpu_om = np.abs(gpu_number-om_number).sum()
    print("the total error_of pixel values between pb and om:",total_error_pb_om)
    print("the total error of pixel values between gpu and om:",total_error_gpu_om)

if __name__ == "__main__" :
    main(parse_arguments(sys.argv[1:]))